package tree;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Stack;
import java.util.concurrent.Delayed;

import parser.RTEXmlParser;

import matching.Matcher;

import data.*;

public class TreeEditDistance extends Matcher {

	private int DELETION_COST = 1;
	private final int SUBSTITUTION_COST = 1;
	private float INSERTION_COST = 1;
	private boolean useIDF = false;
	private boolean useTestTrees = false;
	
	public TreeEditDistance(String task, double threshold) {
		super(task, threshold);
	}

	public void setDelCost(int value) {
		DELETION_COST = value;
	}
	
	public void setUseIDF(boolean value) {
		this.useIDF = value;
	}
	
	public void setUseTestTrees(boolean value) {
		this.useTestTrees = value;
	}
	
	@Override
	public void match() {
		ArrayList<RTEPreprocessedData> dataList = null;
		if(useTestTrees) {
			dataList = new ArrayList<RTEPreprocessedData>();
			dataList.add(generateTestTrees());
		}
		else {
			dataList = RTEPreprocessedData.rtePreprocessedDataList;
		}
		
		numberOfTextCases = dataList.size();
		
		for (RTEPreprocessedData data : dataList) {
			ArrayList<RTENode> t1Roots = new ArrayList<RTENode>();
			ArrayList<RTENode> t2Roots = new ArrayList<RTENode>();

			//add all text sentences to the list of roots
			for (RTESentence rteSentence : data.getTextSentences()) {
				ArrayList<RTENode> sentenceRoots = buildTree(rteSentence.getNodes());
				for (RTENode rootNode : sentenceRoots) {
					t1Roots.add(rootNode);
				}
			}
			
			for (RTESentence rteSentence : data.getHypothesisSentences()) {
				ArrayList<RTENode> sentenceRoots = buildTree(rteSentence.getNodes());
				for (RTENode rootNode : sentenceRoots) {
					t2Roots.add(rootNode);
				}
			}

			RTENode t1Root = mergeForest(t1Roots);
			RTENode t2Root = mergeForest(t2Roots);
						
			ArrayList<RTENode> t1PostOrder = new ArrayList<RTENode>();
			ArrayList<RTENode> t2PostOrder = new ArrayList<RTENode>();
	
			//calculate the post-order ordering of the nodes
			postOrder(t1Root, t1PostOrder);
			postOrder(t2Root, t2PostOrder);
			
			//TODO: can this be done in postOrder method?
			//set the id of the nodes according to their position in the postorder traversal
			for(int i = 0; i < t1PostOrder.size(); i++) {
				t1PostOrder.get(i).setModifiedId(i+1);
			}
			for(int i = 0; i < t2PostOrder.size(); i++) {
				t2PostOrder.get(i).setModifiedId(i+1);
			}
			
			//calculate the left-most decendants of the tree in the post-order ordering
			ArrayList<RTENode> t1LeftMostDecendants = leftMostLeafDecendants(t1PostOrder);
			ArrayList<RTENode> t2LeftMostDecendants = leftMostLeafDecendants(t2PostOrder);
			
			//calculate the key root indicies of the tree
			ArrayList<RTENode> t1KeyRootIndices = calculateKeyRoots(t1Root);
			ArrayList<RTENode> t2KeyRootIndices = calculateKeyRoots(t2Root);
			
//			printDebug(r1, postOrderT1, leftMostDecendantsT1, keyRootIndicesT1);
//			printDebug(r2, postOrderT2, leftMostDecendantsT2, keyRootIndicesT2);
			
			float[][] treeDistance = new float[t1PostOrder.size()][];
			for(int i = 0; i < t1PostOrder.size(); i++) {
				treeDistance[i] = new float[t2PostOrder.size()];
			}
			
			//Main loop
			for(int i = 0; i < t1KeyRootIndices.size(); i++) {
				int t1KeyRoot = t1KeyRootIndices.get(i).getModifiedId();
				for(int j = 0; j < t2KeyRootIndices.size(); j++)  {
					int t2KeyRoot = t2KeyRootIndices.get(j).getModifiedId();
					computeForestDistance(t1KeyRoot, t2KeyRoot, t1LeftMostDecendants, t2LeftMostDecendants, treeDistance, t1PostOrder, t2PostOrder);
				}
			}

			float editDistance = treeDistance[treeDistance.length - 1][treeDistance[0].length - 1];
			
			float idfScore = 0;
			for (RTENode node : t2PostOrder) {
				idfScore += idfPreProcessed.get(node.getLemma());
			}
			float score = 0;
			
			if(useIDF)
				score = editDistance / idfScore;
			else 
				score = (t2PostOrder.size() - editDistance) / (float)t2PostOrder.size();
			
			boolean entails = false;
			if(score >= threshold) { 
				entails = true;
			}	
			this.addJudgement(data.getId(), entails);
			
			if(entails == data.isEntailment()) { 
				entailments++;
			}
			else
				errors.add(data.getId());
			
			if(DELETION_COST > 0)
				System.out.print(data.getId() + ": Edit distance = " + treeDistance[treeDistance.length -1][treeDistance[0].length - 1] + "\n");
//				System.out.println("score = " + score + " guess = " + entails + ". Answer = " + data.isEntailment());
			
		}

		correctness = (float)entailments/(float)numberOfTextCases;
		if(DELETION_COST == 0)	
		System.out.println(this);
		
		RTEXmlParser.WritePredictions("TreeEditDistance", this.getJudgementData());
	}
	
	
	/**
	 * Merges a forest of trees to one tree
	 * @param roots 
	 * @return root node for the new tree
	 */
	private RTENode mergeForest(ArrayList<RTENode> roots) {
		RTENode root = new RTENode("0", "", "root");
		for (RTENode rteNode : roots) {
			root.addChild(rteNode);
		}
		return root;
	}


	/**
	 * Calculates the tree edit distance of two sub-trees
	 * @param t1KeyRoot key root of tree one
	 * @param t2KeyRoot key root of tree two
	 * @param t1Lmld a list of the left most descendant nodes for tree one
	 * @param t2Lmld a list of the left most descendant nodes for tree two
	 * @param td the final tree edit distance matrix
	 * @param t1PostOrder the post ordering of tree one
	 * @param t2PostOrder the post ordering of tree two
	 */
	private void computeForestDistance(int t1KeyRoot, int t2KeyRoot, ArrayList<RTENode> t1Lmld, ArrayList<RTENode> t2Lmld, float[][] td, ArrayList<RTENode> t1PostOrder, ArrayList<RTENode> t2PostOrder){
		//keyroot -1 p.g.a 0 indeksering
 		int l1 = t1Lmld.get(t1KeyRoot -1).getModifiedId() - 1;
		int l2 = t2Lmld.get(t2KeyRoot -1).getModifiedId() - 1;
		
		int lengthI = t1KeyRoot - l1 + 1;
		int lengthJ = t2KeyRoot - l2 + 1;
		
		//initialize forest distance matrix
		float[][] forestDistance = new float[lengthI][];
		for (int j = 0; j < forestDistance.length; j++) {
			forestDistance[j] = new float[lengthJ];
		}
		
		for (int i = 1; i < lengthI; i++) {
			forestDistance[i][0] = forestDistance[i-1][0] + DELETION_COST;
		}
		
		for (int j = 1; j < lengthJ; j++) {
			forestDistance[0][j] = forestDistance[0][j-1] + INSERTION_COST;
		}	

		int x = 1;
				
		for(int i = l1; i < t1KeyRoot; i++) {
			int y = 1;
			for(int j = l2; j < t2KeyRoot; j++) {
				String lemmaI = t1PostOrder.get(i).getLemma();
				String lemmaJ = t2PostOrder.get(j).getLemma();
				
				INSERTION_COST = useIDF ? idfPreProcessed.get(lemmaJ) : 1;
//				System.out.println("INS FOR WORD " + lemmaJ + " is " + INSERTION_COST);
				
				if(t1Lmld.get(i).getModifiedId() == l1+1 && t2Lmld.get(j).getModifiedId() == l2+1) {		
					float delCost = forestDistance[x - 1][y] + DELETION_COST;
					float insCost = forestDistance[x][y - 1] + INSERTION_COST;
					float subCost = forestDistance[x - 1][y - 1] + (lemmaI.equals(lemmaJ) ? 0 : SUBSTITUTION_COST);
					
					forestDistance[x][y] = Math.min(Math.min(delCost, insCost), subCost);					
//					System.out.println("MIN OF: " + forestDistance[x-1][y] + "+1, " + forestDistance[x][y-1] + "+" + INSERTION_COST + ", " + forestDistance[x-1][y-1] + "+1");
//					System.out.println("result = " + forestDistance[x][y]);
					td[i][j] = forestDistance[x][y];
				}
				else {			
					float delCost = forestDistance[x - 1][y] + DELETION_COST;
					float insCost = forestDistance[x][y - 1] + INSERTION_COST;
					float subCost = forestDistance[x-1][y-1] + td[i][j];
					forestDistance[x][y] = Math.min(Math.min(delCost, insCost), subCost);					
				}
				y++;
			}
			x++;
		}
		
//		printForestDistance(forestDistance);
//		printForestDistance(td);
	}
	
	private void printDebug(RTENode root, ArrayList<RTENode> postOrdering, ArrayList<RTENode> leftMostDescendants, ArrayList<RTENode> keyRoots) {
		System.out.println("T1: ");
		PrintTree(root);
		
		System.out.print("PostOrder: ");
		for (RTENode rteNode : postOrdering) {
			System.out.print(rteNode.getModifiedId() + ", ");
		}
		System.out.println();
		
		System.out.print("Left-most decendants: ");
		for (RTENode rteNode : leftMostDescendants) {
			System.out.print(rteNode.getModifiedId() + ", ");
		}
		System.out.println();

		System.out.print("Key root indices: ");
		for (RTENode rteNode : keyRoots) {
			System.out.print(rteNode.getModifiedId() + ", ");
		}
		System.out.println();
	}

	private void printForestDistance(float[][] forest) {
		for (int i = 0; i < forest.length; i++) {
			System.out.print("[ ");
			for (int j = 0; j < forest[0].length; j++) {
				System.out.print(forest[i][j] + " ");
			}
			System.out.println("]");
		}
		System.out.println();
	}

	private ArrayList<RTENode> buildTree(ArrayList<RTENode> nodes) {		
		//loop over all nodes
		ArrayList<RTENode> rootNodes = new ArrayList<RTENode>();
		for (int i = 0; i < nodes.size(); i++) {
			RTENode current = nodes.get(i);
			if(current.getRelationId().length() == 0)
			{
				rootNodes.add(current);
				continue;
			}
			RTENode parent = RTENode.findNode(nodes, current.getRelationId());
			RTENode child = RTENode.findNode(nodes, current.getId());
			parent.addChild(child);	
		}
		return rootNodes;
	}
	
	private void PrintTree(RTENode root) {
		Queue<RTENode> queue = new LinkedList<RTENode>();
		
		queue.add(root);
		
		while(!queue.isEmpty()) {
			RTENode node = queue.poll();
			System.out.print(node.getModifiedId() + "[");
			
			for (RTENode rteNode : node.getChildren()) {
				System.out.print(" "+rteNode.getModifiedId());
				queue.add(rteNode);
			}
			System.out.println(" ]");
		}
	}

	/**
	 * Calculates the post order ordering of the nodes in the tree
	 * @param node
	 * @param nodes 
	 */
	private void postOrder(RTENode node, ArrayList<RTENode> nodes) {		
		for(RTENode n : node.getChildren()) {
			postOrder(n, nodes);
		}
		nodes.add(node);		
	}

	/**
	 * Finds the left most leaf descendants of all nodes in a tree
	 * @param nodes
	 * @return
	 */
	private ArrayList<RTENode> leftMostLeafDecendants(ArrayList<RTENode> nodes) {
		ArrayList<RTENode> lmldNodes = new ArrayList<RTENode>();
		for (RTENode rteNode : nodes) {
			lmldNodes.add(lmld(rteNode));
		}
		return lmldNodes;
	}
	
	private RTENode lmld(RTENode node) {
		if(node.getChildren().size() == 0)
			return node;
		return lmld(node.getLeftChild());
	}
	
	private ArrayList<RTENode> calculateKeyRoots(RTENode root) {
		ArrayList<RTENode> nodes = new ArrayList<RTENode>();
		
		keyRoot(root, nodes);
		nodes.add(root);
		return nodes;
	}

	private void keyRoot(RTENode node, ArrayList<RTENode> nodes) {
		if(node.getChildren().size() == 0)
			return;
		for(RTENode child : node.getChildren()) {
			keyRoot(child, nodes);
		}
		
		if(node.getChildren().size() > 1) {
			for (int i = 1; i < node.getChildren().size(); i++) {
				nodes.add(node.getChildren().get(i));
			}
		}
	}
	
	private RTEPreprocessedData generateTestTrees() {
		ArrayList<RTENode> text = new ArrayList<RTENode>();
		text.add(new RTENode("1", "4", "a"));
		text.add(new RTENode("2", "3", "b"));
		text.add(new RTENode("3", "4", "c"));
		text.add(new RTENode("4", "6", "d"));
		text.add(new RTENode("5", "6", "e"));
		text.add(new RTENode("6", "", "f"));
		
		ArrayList<RTENode> text2 = new ArrayList<RTENode>();
		text2.add(new RTENode("1", "3", "a"));
		text2.add(new RTENode("2", "3", "b"));
		text2.add(new RTENode("3", "4", "d"));
		text2.add(new RTENode("4", "6", "c"));
		text2.add(new RTENode("5", "6", "e"));
		text2.add(new RTENode("6", "", "f"));
		
		ArrayList<RTESentence> textSentence = new ArrayList<RTESentence>();
		textSentence.add(new RTESentence(text));
		
		ArrayList<RTESentence> hypSentence = new ArrayList<RTESentence>();
		hypSentence.add(new RTESentence(text2));
		
		return new RTEPreprocessedData(1337, true, textSentence, hypSentence);
	}
}